#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
高级内存保护系统 - 实现内存保护、防止内存转储和代码注入
通过多种技术保护应用程序内存免受逆向工程和分析
"""
import os
import sys
import time
import platform
import ctypes
import threading
import random
import hashlib
import re
import mmap
from typing import Dict, List, Tuple, Optional, Any, Union, Callable

# 导入项目模块
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))  # noqa
from utils import Logger

class MemoryProtector:
    """高级内存保护系统核心类"""
    def __init__(self):
        """初始化内存保护系统"""
        # 内存保护配置
        self.memory_config = {
            "enable_memory_encryption": True,
            "enable_self_modifying_code": False,  # 自修改代码功能
            "enable_memory_checksumming": True,   # 内存校验和
            "enable_memory_randomization": True,  # 内存随机化
            "enable_integrity_check": True,       # 完整性检查
            "detection_level": "medium",         # low, medium, high
            "response_strategy": "warn",         # log, warn, exit
            "check_interval": 2.0,                # 秒
            "secure_sections": []                 # 需要保护的内存区域
        }
        
        # 线程锁
        self.lock = threading.Lock()
        
        # 监控线程
        self.monitor_thread = None
        self._stop_monitor = False
        
        # 系统类型
        self.system_type = platform.system()
        
        # 内存块和校验和映射
        self.memory_checksums = {}
        
        # 自修改代码区域
        self.self_modifying_sections = {}
        
        # 初始化内存保护
        self._init_platform_specific()
        
        Logger.log_info("高级内存保护系统初始化完成")
    
    def _init_platform_specific(self) -> None:
        """初始化平台特定的内存保护机制"""
        try:
            if self.system_type == "Windows":
                # Windows平台特定初始化
                self._windows_init()
            elif self.system_type == "Linux":
                # Linux平台特定初始化
                self._linux_init()
            elif self.system_type == "Darwin":  # macOS
                # macOS平台特定初始化
                self._macos_init()
        except Exception as e:
            Logger.log_error(f"平台特定初始化失败: {str(e)}")
    
    def _windows_init(self) -> None:
        """Windows平台特定初始化"""
        try:
            # 加载Windows API
            self.kernel32 = ctypes.windll.kernel32
            
            # 定义内存保护常量
            self.PAGE_EXECUTE_READWRITE = 0x40
            self.PAGE_READWRITE = 0x04
            self.PAGE_NOACCESS = 0x01
            self.PAGE_EXECUTE = 0x10
            
            Logger.log_info("Windows平台内存保护初始化完成")
        except Exception as e:
            Logger.log_error(f"Windows平台初始化失败: {str(e)}")
    
    def _linux_init(self) -> None:
        """Linux平台特定初始化"""
        try:
            # Linux平台内存保护初始化
            self.PROT_READ = 1
            self.PROT_WRITE = 2
            self.PROT_EXEC = 4
            self.PROT_NONE = 0
            
            # 尝试打开/proc/self/maps文件以检查内存映射
            try:
                with open("/proc/self/maps", "r") as f:
                    pass
            except Exception as e:
                Logger.log_warning(f"无法访问/proc/self/maps: {str(e)}")
            
            Logger.log_info("Linux平台内存保护初始化完成")
        except Exception as e:
            Logger.log_error(f"Linux平台初始化失败: {str(e)}")
    
    def _macos_init(self) -> None:
        """macOS平台特定初始化"""
        try:
            # macOS平台内存保护初始化
            Logger.log_info("macOS平台内存保护初始化完成")
        except Exception as e:
            Logger.log_error(f"macOS平台初始化失败: {str(e)}")
    
    def start_protection(self) -> bool:
        """启动内存保护"""
        try:
            with self.lock:
                # 检查是否启用内存保护
                if not self.memory_config["enable_memory_encryption"] and \
                   not self.memory_config["enable_memory_checksumming"] and \
                   not self.memory_config["enable_memory_randomization"] and \
                   not self.memory_config["enable_integrity_check"]:
                    Logger.log_info("所有内存保护功能均已禁用，跳过保护启动")
                    return False
                
                # 停止已有的监控线程
                if self.monitor_thread and self.monitor_thread.is_alive():
                    self.stop_protection()
                
                # 重置停止标志
                self._stop_monitor = False
                
                # 创建并启动监控线程
                self.monitor_thread = threading.Thread(target=self._monitor_memory, daemon=True)
                self.monitor_thread.start()
                
                Logger.log_info("内存保护已启动")
                return True
        except Exception as e:
            Logger.log_error(f"启动内存保护失败: {str(e)}")
            return False
    
    def stop_protection(self) -> bool:
        """停止内存保护"""
        try:
            with self.lock:
                # 设置停止标志
                self._stop_monitor = True
                
                # 等待监控线程结束
                if self.monitor_thread and self.monitor_thread.is_alive():
                    self.monitor_thread.join(timeout=3.0)
                    
                    # 如果线程仍在运行，记录警告
                    if self.monitor_thread.is_alive():
                        Logger.log_warning("内存保护监控线程未能正常终止")
                        return False
                
                # 清除内存块映射
                self.memory_checksums.clear()
                
                Logger.log_info("内存保护已停止")
                return True
        except Exception as e:
            Logger.log_error(f"停止内存保护失败: {str(e)}")
            return False
    
    def _monitor_memory(self) -> None:
        """监控内存状态，定期执行内存检查"""
        try:
            while not self._stop_monitor:
                # 执行内存校验和检查
                if self.memory_config["enable_memory_checksumming"]:
                    self._check_memory_checksums()
                
                # 执行内存完整性检查
                if self.memory_config["enable_integrity_check"]:
                    self._check_memory_integrity()
                
                # 等待下一次检查
                wait_time = self.memory_config["check_interval"] + random.uniform(-0.2, 0.2)  # 添加随机波动
                start_time = time.time()
                while time.time() - start_time < wait_time and not self._stop_monitor:
                    time.sleep(0.1)  # 小间隔检查是否需要停止
        except Exception as e:
            Logger.log_error(f"内存监控循环发生错误: {str(e)}")
    
    def protect_data(self, data: Any) -> Any:
        """保护敏感数据"""
        try:
            if not self.memory_config["enable_memory_encryption"]:
                return data
            
            # 根据数据类型选择不同的保护方法
            if isinstance(data, str):
                return self._encrypt_string(data)
            elif isinstance(data, (int, float, bool)):
                return self._obfuscate_value(data)
            elif isinstance(data, (list, tuple, set)):
                return self._protect_collection(data)
            elif isinstance(data, dict):
                return self._protect_dictionary(data)
            else:
                # 对于不支持的类型，直接返回原数据
                return data
        except Exception as e:
            self.logger.log_error(f"数据保护失败: {str(e)}")
            return data
    
    def unprotect_data(self, protected_data: Any) -> Any:
        """解除数据保护"""
        try:
            if not self.memory_config["enable_memory_encryption"]:
                return protected_data
            
            # 根据数据类型选择不同的解除保护方法
            if isinstance(protected_data, str) and self._is_encrypted_string(protected_data):
                return self._decrypt_string(protected_data)
            elif isinstance(protected_data, (int, float, bool)):
                return self._deobfuscate_value(protected_data)
            elif isinstance(protected_data, (list, tuple, set)):
                return self._unprotect_collection(protected_data)
            elif isinstance(protected_data, dict):
                return self._unprotect_dictionary(protected_data)
            else:
                # 对于不支持的类型，直接返回原数据
                return protected_data
        except Exception as e:
            self.logger.log_error(f"数据解除保护失败: {str(e)}")
            return protected_data
    
    def _encrypt_string(self, s: str) -> str:
        """加密字符串"""
        try:
            # 简单的XOR加密算法示例
            # 真实应用中应该使用更强的加密算法
            key = self._get_encryption_key()
            encrypted_bytes = bytes([ord(c) ^ key[i % len(key)] for i, c in enumerate(s)])
            # 转换为Base64以保持可打印性
            import base64
            return f"__ENC__{base64.b64encode(encrypted_bytes).decode('utf-8')}"
        except Exception as e:
            self.logger.log_error(f"字符串加密失败: {str(e)}")
            return s
    
    def _decrypt_string(self, s: str) -> str:
        """解密字符串"""
        try:
            # 检查是否是我们加密的字符串
            if not s.startswith("__ENC__"):
                return s
            
            # 提取加密数据
            import base64
            encrypted_data = s[6:]  # 去掉前缀__ENC__
            encrypted_bytes = base64.b64decode(encrypted_data)
            
            # 解密
            key = self._get_encryption_key()
            decrypted_bytes = bytes([b ^ key[i % len(key)] for i, b in enumerate(encrypted_bytes)])
            return decrypted_bytes.decode('utf-8', errors='replace')
        except Exception as e:
            self.logger.log_error(f"字符串解密失败: {str(e)}")
            return s
    
    def _is_encrypted_string(self, s: str) -> bool:
        """检查字符串是否已加密"""
        return isinstance(s, str) and s.startswith("__ENC__")
    
    def _obfuscate_value(self, value: Union[int, float, bool]) -> Union[int, float, bool]:
        """混淆数值"""
        try:
            if isinstance(value, bool):
                # 布尔值混淆
                # 简单地使用XOR操作
                return value ^ (self._get_obfuscation_seed() % 2 == 0)
            elif isinstance(value, int):
                # 整数混淆
                seed = self._get_obfuscation_seed()
                # 使用一些数学运算来混淆整数
                obfuscated = (value ^ seed) + seed // 2
                return obfuscated
            elif isinstance(value, float):
                # 浮点数混淆
                seed = self._get_obfuscation_seed()
                # 使用一些数学运算来混淆浮点数
                obfuscated = (value * (1 + seed / 1000)) ^ (seed % 100) / 100
                return obfuscated
            return value
        except Exception as e:
            self.logger.log_error(f"数值混淆失败: {str(e)}")
            return value
    
    def _deobfuscate_value(self, value: Union[int, float, bool]) -> Union[int, float, bool]:
        """解除数值混淆"""
        try:
            if isinstance(value, bool):
                # 布尔值解除混淆
                return value ^ (self._get_obfuscation_seed() % 2 == 0)
            elif isinstance(value, int):
                # 整数解除混淆
                seed = self._get_obfuscation_seed()
                # 逆向混淆运算
                deobfuscated = (value - seed // 2) ^ seed
                return deobfuscated
            elif isinstance(value, float):
                # 浮点数解除混淆
                seed = self._get_obfuscation_seed()
                # 逆向混淆运算
                temp = value ^ (seed % 100) / 100
                deobfuscated = temp / (1 + seed / 1000)
                return deobfuscated
            return value
        except Exception as e:
            self.logger.log_error(f"数值解除混淆失败: {str(e)}")
            return value
    
    def _protect_collection(self, collection: Union[list, tuple, set]) -> Union[list, tuple, set]:
        """保护集合类型数据"""
        try:
            # 保护集合中的每个元素
            protected_items = [self.protect_data(item) for item in collection]
            
            # 保持原始数据类型
            if isinstance(collection, list):
                return protected_items
            elif isinstance(collection, tuple):
                return tuple(protected_items)
            elif isinstance(collection, set):
                return set(protected_items)
            return collection
        except Exception as e:
            self.logger.log_error(f"集合数据保护失败: {str(e)}")
            return collection
    
    def _unprotect_collection(self, collection: Union[list, tuple, set]) -> Union[list, tuple, set]:
        """解除集合类型数据保护"""
        try:
            # 解除集合中每个元素的保护
            unprotected_items = [self.unprotect_data(item) for item in collection]
            
            # 保持原始数据类型
            if isinstance(collection, list):
                return unprotected_items
            elif isinstance(collection, tuple):
                return tuple(unprotected_items)
            elif isinstance(collection, set):
                return set(unprotected_items)
            return collection
        except Exception as e:
            self.logger.log_error(f"集合数据解除保护失败: {str(e)}")
            return collection
    
    def _protect_dictionary(self, dictionary: dict) -> dict:
        """保护字典类型数据"""
        try:
            # 保护字典中的键值对
            protected_dict = {self.protect_data(k): self.protect_data(v) for k, v in dictionary.items()}
            return protected_dict
        except Exception as e:
            self.logger.log_error(f"字典数据保护失败: {str(e)}")
            return dictionary
    
    def _unprotect_dictionary(self, dictionary: dict) -> dict:
        """解除字典类型数据保护"""
        try:
            # 解除字典中键值对的保护
            unprotected_dict = {self.unprotect_data(k): self.unprotect_data(v) for k, v in dictionary.items()}
            return unprotected_dict
        except Exception as e:
            self.logger.log_error(f"字典数据解除保护失败: {str(e)}")
            return dictionary
    
    def add_secure_section(self, section_id: str, data: Any, protection_level: str = "medium") -> bool:
        """添加安全内存区域"""
        try:
            with self.lock:
                # 保护数据
                protected_data = self.protect_data(data)
                
                # 计算数据的校验和
                checksum = self._calculate_checksum(protected_data)
                
                # 存储保护后的数据和校验和
                self.secure_sections[section_id] = {
                    "data": protected_data,
                    "checksum": checksum,
                    "protection_level": protection_level,
                    "last_check": time.time()
                }
                
                self.logger.log_info(f"添加安全内存区域: {section_id}")
                return True
        except Exception as e:
            self.logger.log_error(f"添加安全内存区域失败: {str(e)}")
            return False
    
    def get_secure_section(self, section_id: str) -> Any:
        """获取安全内存区域的数据"""
        try:
            with self.lock:
                if section_id not in self.secure_sections:
                    self.logger.log_error(f"安全内存区域不存在: {section_id}")
                    return None
                
                # 验证数据完整性
                secure_section = self.secure_sections[section_id]
                current_checksum = self._calculate_checksum(secure_section["data"])
                
                if current_checksum != secure_section["checksum"]:
                    self.logger.log_warning(f"安全内存区域数据已被篡改: {section_id}")
                    self._handle_memory_breach("数据篡改", section_id)
                    return None
                
                # 更新最后检查时间
                secure_section["last_check"] = time.time()
                
                # 解除数据保护
                return self.unprotect_data(secure_section["data"])
        except Exception as e:
            self.logger.log_error(f"获取安全内存区域数据失败: {str(e)}")
            return None
    
    def remove_secure_section(self, section_id: str) -> bool:
        """移除安全内存区域"""
        try:
            with self.lock:
                if section_id in self.secure_sections:
                    del self.secure_sections[section_id]
                    self.logger.log_info(f"移除安全内存区域: {section_id}")
                    return True
                return False
        except Exception as e:
            self.logger.log_error(f"移除安全内存区域失败: {str(e)}")
            return False
    
    def _check_memory_checksums(self) -> None:
        """检查内存校验和"""
        try:
            with self.lock:
                for section_id, section_info in list(self.secure_sections.items()):
                    # 计算当前数据的校验和
                    current_checksum = self._calculate_checksum(section_info["data"])
                    
                    # 比较校验和
                    if current_checksum != section_info["checksum"]:
                        self.logger.log_warning(f"检测到内存篡改: {section_id}")
                        self._handle_memory_breach("校验和不匹配", section_id)
                        
                        # 对于高保护级别区域，自动清理数据
                        if section_info["protection_level"] == "high":
                            self.logger.log_info(f"自动清理高风险内存区域: {section_id}")
                            del self.secure_sections[section_id]
        except Exception as e:
            self.logger.log_error(f"内存校验和检查失败: {str(e)}")
    
    def _check_memory_integrity(self) -> None:
        """检查内存完整性"""
        try:
            # 根据操作系统类型执行不同的内存完整性检查
            if self.system_type == "Windows":
                self._windows_check_memory_integrity()
            elif self.system_type == "Linux":
                self._linux_check_memory_integrity()
            elif self.system_type == "Darwin":  # macOS
                self._macos_check_memory_integrity()
        except Exception as e:
            self.logger.log_error(f"内存完整性检查失败: {str(e)}")
    
    def _windows_check_memory_integrity(self) -> None:
        """Windows平台内存完整性检查"""
        try:
            # 检查是否有进程打开了当前进程的内存
            # 这是一个简化的实现，实际应用中可能需要更复杂的方法
            if self.memory_config["detection_level"] == "high":
                # 获取当前进程句柄
                process_handle = self.kernel32.OpenProcess(0x0400, False, os.getpid())
                
                if process_handle:
                    # 这里可以添加更复杂的内存完整性检查
                    # 例如，检查进程的内存映射、保护标志等
                    self.kernel32.CloseHandle(process_handle)
        except Exception as e:
            self.logger.log_error(f"Windows内存完整性检查失败: {str(e)}")
    
    def _linux_check_memory_integrity(self) -> None:
        """Linux平台内存完整性检查"""
        try:
            # 检查/proc/self/maps文件以监控内存映射变化
            if self.memory_config["detection_level"] in ["medium", "high"]:
                try:
                    with open("/proc/self/maps", "r") as f:
                        maps_content = f.read()
                        
                        # 检查是否有可疑的内存映射
                        # 例如，具有可执行权限的匿名映射可能是代码注入
                        executable_anonymous_maps = re.findall(r'([0-9a-f]+)-([0-9a-f]+)\s+rwxp.*\s+00:00\s+0', maps_content)
                        
                        if executable_anonymous_maps and len(executable_anonymous_maps) > 5:  # 假设正常情况下不会有太多可执行匿名映射
                            self.logger.log_warning(f"检测到可疑的可执行匿名内存映射: {len(executable_anonymous_maps)}")
                            self._handle_memory_breach("可疑内存映射", "linux_maps")
                except Exception as e:
                    self.logger.log_error(f"读取/proc/self/maps失败: {str(e)}")
        except Exception as e:
            self.logger.log_error(f"Linux内存完整性检查失败: {str(e)}")
    
    def _macos_check_memory_integrity(self) -> None:
        """macOS平台内存完整性检查"""
        try:
            # macOS平台内存完整性检查
            # 这是一个简化的实现，实际应用中可能需要更复杂的方法
            self.logger.log_info("macOS内存完整性检查执行")
        except Exception as e:
            self.logger.log_error(f"macOS内存完整性检查失败: {str(e)}")
    
    def _handle_memory_breach(self, breach_type: str, section_id: str) -> None:
        """处理内存安全漏洞"""
        strategy = self.memory_config["response_strategy"]
        
        if strategy == "log":
            # 仅记录日志
            self.logger.log_info(f"内存安全漏洞: {breach_type} (区域: {section_id}) - 仅记录日志")
        elif strategy == "warn":
            # 记录警告
            self.logger.log_warning(f"内存安全漏洞: {breach_type} (区域: {section_id}) - 发出警告")
        elif strategy == "exit":
            # 记录错误并退出应用程序
            self.logger.log_error(f"内存安全漏洞: {breach_type} (区域: {section_id}) - 执行退出策略")
            # 优雅地清理资源
            self.stop_protection()
            # 退出应用程序
            sys.exit(1)
        else:
            # 默认策略：记录警告
            self.logger.log_warning(f"内存安全漏洞: {breach_type} (区域: {section_id}) - 使用默认策略")
    
    def _calculate_checksum(self, data: Any) -> str:
        """计算数据的校验和"""
        try:
            # 将数据序列化为字节
            serialized_data = self._serialize_data(data)
            
            # 计算SHA-256哈希值
            return hashlib.sha256(serialized_data).hexdigest()
        except Exception as e:
            self.logger.log_error(f"计算校验和失败: {str(e)}")
            return ""
    
    def _serialize_data(self, data: Any) -> bytes:
        """序列化数据为字节"""
        try:
            import pickle
            return pickle.dumps(data)
        except Exception as e:
            self.logger.log_error(f"数据序列化失败: {str(e)}")
            return b""
    
    def _get_encryption_key(self) -> bytes:
        """获取加密密钥"""
        # 在实际应用中，应该使用更安全的方式存储和获取密钥
        # 这里使用进程ID和当前时间作为密钥源，仅作为示例
        pid = os.getpid()
        current_time = int(time.time() * 1000)
        key_source = f"{pid}_{current_time}_{platform.node()}"
        return hashlib.sha256(key_source.encode()).digest()[:16]  # 使用16字节密钥
    
    def _get_obfuscation_seed(self) -> int:
        """获取混淆种子"""
        # 在实际应用中，应该使用更安全的方式生成混淆种子
        # 这里使用进程ID和当前时间作为种子源，仅作为示例
        pid = os.getpid()
        current_time = int(time.time() * 1000)
        seed_source = f"{pid}_{current_time}_{random.randint(0, 1000)}"
        return int(hashlib.sha256(seed_source.encode()).hexdigest(), 16) % 1000000
    
    def set_config(self, config: Dict) -> bool:
        """设置内存保护配置"""
        try:
            with self.lock:
                # 验证配置项
                for key, value in config.items():
                    if key in self.memory_config:
                        # 根据配置项类型进行验证
                        if key == "detection_level" and value not in ["low", "medium", "high"]:
                            self.logger.log_error(f"无效的检测级别: {value}")
                            continue
                        if key == "response_strategy" and value not in ["log", "warn", "exit"]:
                            self.logger.log_error(f"无效的响应策略: {value}")
                            continue
                        if key in ["check_interval"] and not isinstance(value, (int, float)):
                            self.logger.log_error(f"无效的时间值: {value}")
                            continue
                        
                        # 更新配置
                        self.memory_config[key] = value
                        self.logger.log_info(f"更新内存保护配置: {key} = {value}")
                    else:
                        self.logger.log_warning(f"未知的配置项: {key}")
            
            return True
        except Exception as e:
            self.logger.log_error(f"设置内存保护配置失败: {str(e)}")
            return False
    
    def get_memory_stats(self) -> Dict:
        """获取内存保护统计信息"""
        try:
            with self.lock:
                stats = {
                    "secure_sections_count": len(self.secure_sections),
                    "memory_config": self.memory_config.copy(),
                    "is_monitoring": self.monitor_thread is not None and self.monitor_thread.is_alive()
                }
                return stats
        except Exception as e:
            self.logger.log_error(f"获取内存统计信息失败: {str(e)}")
            return {}

if __name__ == "__main__":
    # 示例用法
    memory_protector = MemoryProtector()
    
    # 设置配置
    memory_protector.set_config({
        "detection_level": "high",
        "response_strategy": "warn",
        "check_interval": 1.0
    })
    
    # 添加安全内存区域
    sensitive_data = {"api_key": "secret_key_123", "password": "secure_password"}
    memory_protector.add_secure_section("credentials", sensitive_data, "high")
    
    # 启动内存保护
    memory_protector.start_protection()
    
    try:
        # 模拟应用程序运行
        for i in range(10):
            # 获取受保护的数据
            credentials = memory_protector.get_secure_section("credentials")
            print(f"获取到的凭据: {credentials}")
            
            # 保护一些数据
            protected_value = memory_protector.protect_data("敏感信息")
            print(f"保护后的值: {protected_value}")
            
            # 解除保护
            original_value = memory_protector.unprotect_data(protected_value)
            print(f"解除保护后的值: {original_value}")
            
            time.sleep(2)
    except KeyboardInterrupt:
        # 停止内存保护
        memory_protector.stop_protection()
        print("内存保护已停止")